import asyncio
import json
import os
import time
import fitz
import cv2
import numpy as np
from PIL import Image
from channels.generic.websocket import AsyncWebsocketConsumer
from django.conf import settings

from file.models import  ImageVector, TextVector
from core.utils.extract_img import extract_table_img
from core.utils.text_split import TextSplit
from core.model_loader import get_text_embedding_model, get_image_embedding_model, get_table_detect_model


from file.constants.task_enum import TaskStatusEnum
from asgiref.sync import sync_to_async

text_split_obj = TextSplit()
text_embedding_model = get_text_embedding_model()
image_embedding_model = get_image_embedding_model()
table_det_model = get_table_detect_model()


class FileProcessConsumer(AsyncWebsocketConsumer):
    async def connect(self):
        await self.accept()

    async def disconnect(self, code):
        pass

    async def receive(self, text_data):
        json_data = json.loads(text_data)
        task_code = json_data.get("taskCode")
        # full_file_path = os.path.join(settings.MEDIA_ROOT, file_path)
        if task_code != "start":
            return
        data = json_data.get("data")
        file_id = data.get("file_id")
        knowbase_id = data.get("know_base_id")
        file_path = data.get("file_path")
        last_level = os.path.basename(settings.MEDIA_ROOT)
        file_path = os.path.join(last_level, file_path)
        out_dir = os.path.dirname(file_path)
        await self.process_pdf_file(file_path, out_dir,knowbase_id, file_id)

    async def process_pdf_file(self, file_path,  out_dir, knowbase_id, file_id):
        text_callback = lambda message ,status = "pending":asyncio.create_task(
            self.send_progress("textTask",message,status)
        )

        image_callback = lambda message ,status = "pending":asyncio.create_task(
            self.send_progress("imageTask",message,status)
        )

        table_callback = lambda message, status="pending": asyncio.create_task(
            self.send_progress("tableTask", message, status)
        )

        await extract_text_embedding_from_pdf_file(file_path, knowbase_id, file_id,text_callback)
        await  extract_images_from_pdf_file(file_path, out_dir,knowbase_id, file_id, image_callback)
        await extract_tables_from_pdf_file(file_path, out_dir, knowbase_id, file_id,table_callback)

    async def send_progress(self, task_name, message, status):
        await self.send(text_data=json.dumps({
            'task': task_name,
            'message': message,
            'status': status
        }))


async def extract_text_embedding_from_pdf_file(file_path, knowbase_id, file_id, progress_callback):
    start_time_0 = time.time()

    #文本提取
    await(progress_callback("准备进行文本提取...", status = TaskStatusEnum.PENDING.code))
    await(progress_callback("文本提取中...", status=TaskStatusEnum.RUNNING.code))
    pdf_doc = fitz.open(file_path)
    texts = []
    for page in pdf_doc:
        text = page.get_text()
        texts.append(text)
    pdf_doc.close()
    end_time = time.time()
    await(progress_callback(f"文本提取完成,耗时{end_time - start_time_0}s", status=TaskStatusEnum.RUNNING.code))


    #文本分割
    start_time = time.time()
    await(progress_callback("准备进行文本分割...",status = TaskStatusEnum.RUNNING.code))
    await(progress_callback("文本分割中...", status=TaskStatusEnum.RUNNING.code))
    texts = text_split_obj.split_texts(texts)
    texts = text_split_obj.add_meta_data(texts, file_path, knowbase_id, file_id)
    end_time = time.time()
    await(progress_callback(f"文本分割完成,耗时{end_time - start_time}s,共分割为{len(texts)}块文本块", status=TaskStatusEnum.RUNNING.code))
    # for i, text in enumerate(texts):
    #     if "\x00" in text.page_content:
    #         print(f"Found NUL in text at index {i}: {repr(text.page_content[:100])}")


    #张量解析
    start_time = time.time()
    await progress_callback("文本张量解析中...",status = TaskStatusEnum.RUNNING.code)
    if len(texts) != 0:
        texts_embedding = [text_embedding_model.get_text_embedding(text.page_content)
                           for text in texts]
    end_time = time.time()
    await progress_callback(f"文本张量解析完成，耗时{end_time - start_time}s",status =TaskStatusEnum.RUNNING.code)

    #张量数据创建及存储
    start_time = time.time()
    await progress_callback("文本张量存储中...", status=TaskStatusEnum.RUNNING.code)
    objects = [
        TextVector(
            embedding_content=text.page_content.replace("\x00",""),
            embedding_vector=np.array(texts_embedding[idx]),
            file_id=file_id,
            know_base_id=knowbase_id,
        )
        for idx, text in enumerate(texts)
    ]
    await sync_to_async(TextVector.objects.bulk_create)(objects)
    end_time = time.time()
    await progress_callback(f"文本张量存储完成，耗时{end_time - start_time}s", status=TaskStatusEnum.RUNNING.code)
    await progress_callback(f"文本处理完成，共耗时{end_time - start_time_0}s", status=TaskStatusEnum.DONE.code)


async def extract_images_from_pdf_file(file_path, out_dir, knowbase_id, file_id, progress_callback):
    start_time_0 = time.time()

    # 图片提取
    await(progress_callback("准备进行图片提取...", status=TaskStatusEnum.PENDING.code))
    await(progress_callback("图片提取中...", status=TaskStatusEnum.RUNNING.code))
    pdf_doc = fitz.open(file_path)
    out_dir = os.path.join(out_dir, 'images')
    os.makedirs(out_dir, exist_ok=True)
    count = 0
    for p_num in range(len(pdf_doc)):
        page = pdf_doc[p_num]
        imgs = page.get_images(full=True)
        for img_idx, img in enumerate(imgs):
            xref = img[0]
            base_img = pdf_doc.extract_image(xref)
            img_bytes = base_img['image']
            img_ext = base_img['ext']

            img_filename = f"page_{p_num + 1}_image_{img_idx + 1}.{img_ext}"
            img_path = os.path.join(out_dir, img_filename)

            with open(img_path, 'wb') as img_file:
                img_file.write(img_bytes)
            await progress_callback(f"已提取第{p_num + 1}页第{img_idx + 1}个图片",status = TaskStatusEnum.RUNNING.code)
            print('Saved: {}'.format(img_path))
            count += 1
    pdf_doc.close()
    end_time = time.time()
    await progress_callback(f"图片提取完成,耗时{end_time-start_time_0}s,共提取{count}个图片",status = TaskStatusEnum.RUNNING.code)
    await extract_images_embedding_from_dir( out_dir,'image', knowbase_id, file_id, progress_callback)
    end_time = time.time()
    await progress_callback(f"图片处理完成,共耗时{end_time - start_time_0}s,共提取{count}个图片", status=TaskStatusEnum.DONE.code)


async def extract_tables_from_pdf_file(file_path, out_dir, knowbase_id, file_id, progress_callback):
    start_time_0 = time.time()

    #表格提取
    await(progress_callback("准备进行表格提取...", status=TaskStatusEnum.PENDING.code))
    await(progress_callback("表格提取中...", status=TaskStatusEnum.RUNNING.code))
    pdf_doc = fitz.open(file_path)
    out_dir = os.path.join(out_dir, 'tables')
    os.makedirs(out_dir, exist_ok=True)
    count = 0
    for p_num in range(len(pdf_doc)):
        page = pdf_doc[p_num]
        pix = page.get_pixmap(dpi=120, alpha=False)
        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)  # 创建PIL图像对象
        img = np.array(img)
        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        result = table_det_model(img_bgr)
        for table_num, res in enumerate(result):
            file_name = f"page_{p_num + 1}_table_{table_num + 1}.png"
            img_path = os.path.join(out_dir, file_name)
            lt, rt, rb, lb = res["lt"], res["rt"], res["rb"], res["lb"]
            wrapped_img = extract_table_img(img_bgr, lt, rt, rb, lb)
            # success = cv2.imwrite(img_path, wrapped_img)
            # cv2.imwrite不支持中文路径保存
            # 编码为png格式
            success, encoded_img = cv2.imencode('.png', wrapped_img)
            if success:
                with open(img_path, 'wb') as f:
                    f.write(encoded_img.tobytes())
                print('Saved: {}'.format(img_path))
                count += 1
                await progress_callback(f"已提取第{p_num + 1}页第{table_num + 1}个表格")
            else:
                print('Saved失败: {}'.format(img_path))
    pdf_doc.close()
    end_time = time.time()
    await progress_callback(f"表格提取完成,耗时{end_time - start_time_0}s,共提取{count}个表格",
                            status=TaskStatusEnum.RUNNING.code)
    await extract_images_embedding_from_dir(out_dir, 'table',knowbase_id, file_id, progress_callback)
    end_time = time.time()
    await progress_callback(f"表格处理完成,共耗时{end_time - start_time_0}s,共提取{count}个表格",
                            status=TaskStatusEnum.DONE.code)



async def extract_images_embedding_from_dir( out_dir, flag, knowbase_id, file_id, progress_callback):
    start_time = time.time()
    if flag == 'table':
        text_flag = "表格"
    else:
        text_flag = "图片"
    await progress_callback(f"{text_flag}张量解析中...", status=TaskStatusEnum.RUNNING.code)
    image_embeddings = []
    img_files = os.listdir(out_dir)
    for img_file in img_files:
        try:
            image = Image.open(os.path.join(out_dir, img_file))
            image_embedding = image_embedding_model.get_image_embedding(image)
            image_embeddings.append(image_embedding)
        except Exception as e:
            print(e)
    end_time = time.time()
    await progress_callback(f"{text_flag}张量解析完成，耗时{end_time-start_time}", status=TaskStatusEnum.RUNNING.code)

    start_time = time.time()
    await progress_callback(f"{text_flag}张量存储中...", status=TaskStatusEnum.RUNNING.code)
    objects = [
        ImageVector(
            img_description=flag,
            embedding_vector=image_embeddings[idx],
            image_path=os.path.join(out_dir, img_files[idx]),
            know_base_id=knowbase_id,
            file_id=file_id
        )
        for idx in range(len(image_embeddings))
    ]
    await sync_to_async(ImageVector.objects.bulk_create)(objects)
    end_time = time.time()
    await progress_callback(f"{text_flag}张量存储完成，耗时{end_time-start_time}s", status=TaskStatusEnum.RUNNING.code)